# -*- coding: utf-8 -*-
# @Time    : 2024/3/18 17:02
# @Author  : CHNJX
# @File    : telnet_client.py
# @Desc    : telnet客户端
import time

from logging import Logger
import telnetlib


class TelnetClient:

    def __init__(self, host, port, username, password, mac_address=None, timeout=10):
        self.logger = Logger('test')
        self.host = host
        self.port = port
        self.username = username
        self.password = password
        self.timeout = timeout
        self.mac_address = mac_address
        self.telnet = None
        self.connect()
        self.login()

    def login(self):
        """
        登录
        """
        try:
            self.logger.info("等待登录提示...")
            self.telnet.read_until(b"login: ", timeout=20)
            self.telnet.write(self.username.encode('ascii') + b"\n")

            self.logger.info("等待密码提示...")
            self.telnet.read_until(b"Password: ", timeout=20)
            self.telnet.write(self.password.encode('ascii') + b"\n")
        except EOFError:
            self.logger.error("连接被关闭")
            self.connect()  # 尝试重新连接
            self.login()  # 重新登录
        except Exception as e:
            self.logger.error("登录过程中发生错误: %s", e)
            raise

    def send(self, cmd):
        """
        发送命令
        """
        self.logger.debug("发送命令: {}".format(cmd))
        cmd = cmd.replace('\r', '')
        self.telnet.write(cmd.encode('ascii') + b"\n")

        # 假设命令执行完成后的提示符为 '# '
        end_prompt = "# "
        response = ""

        while True:
            chunk = self.telnet.read_very_eager().decode('ascii')
            response += chunk

            # 检查响应是否包含结束标志
            if end_prompt in response:
                # 确保结束标志不在中间，而是作为单独行的开始
                lines = response.splitlines()
                if lines and lines[-1].endswith(end_prompt):
                    return response

            # 如果没有找到结束标志，则继续等待更多数据
            time.sleep(0.1)  # 添加一个短暂的延迟，防止频繁无意义地轮询

    def reconnect(self, max_retry=5):
        """
        重连
        """
        retry = 0  # 当前重试次数
        wait_time = 10  # 等待时间，每次重试之前等待，时间可以根据设备启动时间调整

        while retry < max_retry:
            try:
                self.logger.info("telnet重连尝试 %d/%d", retry + 1, max_retry)
                time.sleep(wait_time)  # 增加延迟，等待设备可能的重启
                self.close()  # 关闭当前连接
                self.connect()  # 尝试重新连接
                self.login()  # 登录
                self.logger.info("telnet重连成功")
                break
            except Exception as e:
                self.logger.error("telnet重连失败: %s", e)
                retry += 1
                wait_time *= 2  # 每次失败后增加等待时间
                if retry == max_retry:
                    self.logger.error("telnet重连尝试失败，将不再重试")
                    raise  # 重试次数用尽，抛出异常

    def connect(self):
        """
        连接telnet
        """
        self.logger.info("telnet连接")
        # 这里优先使用mac地址获取IP地址
        host = self.host
        try:
            self.telnet = telnetlib.Telnet(host, self.port, self.timeout)
        except Exception as e:
            self.logger.error("无法建立连接: %s", e)
            raise

    def close(self):
        """关闭连接"""
        self.logger.info("telnet关闭连接")
        if self.telnet:
            self.telnet.close()
            self.telnet = None  # 重置telnet对象，防止再次使用无效对象

    def read_very_eager(self):
        self.telnet.read_very_eager()

    def read_until(self, expected, timeout=1):
        try:
            data = self.telnet.read_until(expected, timeout)
            return data.decode('ascii', errors='ignore')
        except EOFError:
            self.logger.error("Telnet连接已断开")
            return None
        except Exception as e:
            self.logger.error("读取Telnet数据时发生错误: %s", e)
            return None

    def read_some(self):
        try:
            data = self.telnet.read_some()
            return data.decode('ascii', errors='ignore')
        except EOFError:
            self.logger.error("Telnet连接已断开")
            return None
        except Exception as e:
            self.logger.debug("读取Telnet数据时发生错误: %s", e)
            return None
